perf: stop holding the fair pool lock across blocking memory calls - #5613
perf: stop holding the fair pool lock across blocking memory calls#5613dwsmith1983 wants to merge 34 commits into
Conversation
CometFairMemoryPool held its mutex across the JNI calls into Spark task memory manager, which can block for seconds while Spark spills, so every native thread sharing a task pool serialized behind whichever thread was acquiring. The fairness check and the reservation are now one short locked step, the blocking call runs unlocked, and the reservation rolls back if the JVM fails to back it or the call panics. Fairness semantics are unchanged: concurrent grows still cannot jointly exceed pool_size divided by the consumer count. The JNI boundary moved behind a small trait so the pool finally has tests: fairness rejection, limit tightening on register, partial-grant rollback, panic rollback, a blocking test that took ten seconds on the old code and 30ms now, and an eight-thread stress test.
196d709 to
626bc39
Compare
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Reviewed 626bc395091f313ae93b8867c1cb1c846dc7fc53 against 8729f6e6adf7091e18a48670e790d4ba8fd41e51. I found one P2 in the new concurrent release path, detailed inline.
The focused Spark memory-pool component probe reproduced the missing-task exception. No full Comet native/JNI query suite ran. CI, CodeQL, and Delta Contrib Build Gate currently report action_required, with no test checks recorded.
For this performance change, please include matched BASE/HEAD microbenchmarks with one and multiple native threads, full and partial grants, and consumer registration changes. Report completed operations, grow/release latency, error counts, and final Rust/Spark balances. The parked stub test establishes lock behavior but does not measure production JNI throughput.
| state.used -= subtractive; | ||
| } |
There was a problem hiding this comment.
[P2] Keep waiting Spark acquisitions registered during a full release
Could you handle Spark's same-task wait/release contract before allowing this release to overlap an acquire? With fair_unified, two native consumers can now enter the same CometTaskMemoryManager concurrently. In a 100-unit executor pool, let another task hold 90 and this task hold 10. This task's next 10-unit grow waits below Spark's 1/(2N) minimum. Freeing its last 10 units on another native thread removes its entry from ExecutionMemoryPool.memoryForTask and wakes the grower. The grower then indexes the removed entry and throws NoSuchElementException: key not found. Spark's release bypasses the task monitor held by the waiting acquire, so that monitor does not prevent this interleaving. The Rust provisional reservation does not keep Spark's entry alive.
I reproduced the failure using unchanged Spark 3.5.9 pool source with only logging/annotation/memory-mode scaffolding. A scheduling control modeling the previous serialization completed after the other task freed its memory. The relevant map lifecycle is also present in 4.0.4 source. This was a component probe plus JNI source tracing, not a full Comet query reproduction. Please make full releases safe while grows are pending and add a regression that exercises Spark's memory manager.
There was a problem hiding this comment.
Confirmed against the pool source, thanks for the repro. Blocking the release until the in-flight acquires drain turned out to deadlock in testing, since the parked acquire can be waiting on exactly the memory that release frees. So the fix defers instead: a release that would zero the task's balance while acquires are in flight frees n-1 bytes right away (that is what wakes the waiter) and holds the last byte, which the final completing acquire pays off. At most one byte is ever deferred and it always settles once the acquires finish. The test stub now models the entry lifecycle (created on acquire, removed at zero, a woken waiter fails if the entry is gone) and reproduced this crash before the fix. It also turned out one of our existing tests was exercising the same broken pattern.
There was a problem hiding this comment.
The new deferral fixes the already-pending case, but I can still reproduce the same missing-task failure through a late-arriving acquire on current head eb410e51.
At current lines 252-259, plan_release can see pending_acquires == 0, schedule the whole balance, and drop the state lock before release reaches Spark. A new try_grow can then increment pending_acquires and park in Spark while the old balance is still present. The already-planned release removes the task entry, and the waiter resumes with NoSuchElementException.
I reproduced this with the exact head's production state machine and a gated bridge: hold 10, plan and pause the full release, start and park a grow of 10, then resume the release. The waiter panics with key not found: task entry removed while acquire waited. paying_deferred fences only deferred payments, so could you also coordinate ordinary zeroing releases with newly starting acquires?
There was a problem hiding this comment.
Also covered by db1f1bc. With the anchor there is no zeroing release left to coordinate: once any reservation exists the anchor is held, so a full release leaves the task at one byte and the entry survives.
Two related windows came out of review and are closed in the same commit. A grant that covers the request but not the extra byte is handed back as a short grant rather than running unanchored, and every acquire that starts while the anchor request is still parked carries its own extra byte, with the first full grant keeping it and later ones returning theirs. Your gated repro is pinned as late_acquire_survives_a_release_already_on_its_way, alongside a test for the concurrent in-flight case; both failed on eb410e5 with key not found and pass now. 30 loops each in debug and release are clean.
The earlier build-gate failure compared dylib sizes on a change confined to fair_pool.rs, so it looks like the size check rather than this branch.
There was a problem hiding this comment.
One more window closed in 353541d, found while probing the anchor bootstrap with a third task. Before the anchor lands, a short grant is handed back whole; if another task's entry disappears in that gap (raising this task's minimum share) and a sibling acquire of this task then parks, the rollback release zeroes the entry under it. A small bootstrap mutex now serializes anchor carriers from the bridge acquire through the rollback release. Releases never take it, so a parked carrier cannot starve anyone, and non-carriers only exist once the anchor is held. Pinned as short_grant_rollback_cannot_land_under_a_sibling_parked_in_spark, which failed with key not found before the change.
A release that would zero the JVM-side balance while other acquires are still in flight frees all but one byte immediately and holds the last byte until the in-flight acquires complete. Spark drops the task's accounting entry when its balance hits zero, so a parked acquire waking after that point indexes a missing entry and fails. Blocking the release instead can deadlock because the parked acquire may be waiting for the very memory the release frees. The stub task memory now models the entry lifecycle so the regression is covered.
|
@sunchao any more feedback here? |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Reviewed eb410e5175417256f0223deaee1b122d760ae077 against ef62b46306e925bc51e7d7f29922c1870eb729e7. I found two actionable issues in the revised release/acquire protocol: the held-back byte can leave a valid Spark acquire waiting indefinitely, and an ordinary full release can still be overtaken by a newly starting acquire. The first is inline below. I added the second as a follow-up on the existing unresolved thread because it is the same missing-task failure under a different ordering.
Local validation: the focused native fair-pool suite passed 12/12 tests with 241 filtered out. An exact-head component probe reproduced the ordinary-release ordering failure. A probe using unmodified Spark 4.1.3 ExecutionMemoryPool reproduced both the missing-key failure and the one-byte wait. These are component probes, not a full Comet native/JNI query run.
Current CI has 63 successful checks, 9 skipped checks, and one failed Delta build gate. Its log fails the contrib-enabled-versus-default libcomet size invariant; this PR changes only fair_pool.rs.
| .jvm_held | ||
| .checked_sub(bytes) | ||
| .expect("released more bytes than the JVM side holds"); | ||
| if bytes > 0 && state.jvm_held == 0 && state.pending_acquires > 0 { |
There was a problem hiding this comment.
[P1] Avoid retaining the byte a waiting acquire may need
Could you avoid withholding one byte until finish_acquire? This can deadlock the default off-heap pool. In a 1 GiB Spark execution pool, let another task hold 900 MiB, let this task hold 100 MiB, and let a second native consumer for this task request 100 MiB. Spark parks that request because this task is below its 1/(2N) minimum share. When the holder frees 100 MiB, this branch sends only 100 MiB - 1. On wake, Spark computes toGrant = 100 MiB - 1; because that is short of the request and curMem + toGrant = 100 MiB is still below 256 MiB, it waits again. The deferred byte is paid only by finish_acquire, which cannot run while this acquire is waiting.
I reproduced this against unmodified Spark 4.1.3 ExecutionMemoryPool: retaining one byte left the grower in WAITING, and freeing one additional byte from the other task let it complete. The new stub test misses this because it grants the full request after any release without reapplying Spark's free-memory and minimum-share checks. Could you preserve the task entry without withholding capacity needed by the waiter?
There was a problem hiding this comment.
Fixed in db1f1bc. The deferral is gone: every release now goes to the JVM whole, so a parked acquire sees the full freed amount and Spark's minimum-share check passes. The task entry is kept alive by a permanent anchor instead. The first acquire asks for one extra byte and the pool holds it until it drops, so the balance never reaches zero mid task. That means the task stays in Spark's active-task set for the pool's lifetime and retains one byte, which the header comment now states.
Your scenario is pinned as min_share_wait_is_granted_after_the_holder_frees_its_memory_in_full. The stub now models ExecutionMemoryPool.acquireMemory from 4.1.3 with the per-task entry lifecycle, the 1/N and 1/(2N) shares, the wait loop, and notifyAll on release, with a bounded wait that fails the test instead of hanging. It failed on the previous head with the waiter timing out and passes now.
Spark's ExecutionMemoryPool removes a task's entry when its balance hits zero, and an acquire parked inside Spark indexes that entry on wake. The previous fix held one byte back from a zeroing release while acquires were in flight and paid it back later. That withheld byte starved Spark's minimum share check, which grants a parked request only when the freed bytes cover it in full, so the waiter slept forever; and a release planned before a late acquire parked could still remove the entry. The pool now asks for one extra byte with every acquire that starts before the anchor lands, keeps the first one until the pool drops, and sends every release whole. The task's balance never returns to zero mid task, so both failure modes are impossible by construction. The test double now models ExecutionMemoryPool's entry lifecycle, share rules and wait loop, and covers the minimum share scenario, the late acquire race and the in flight anchor race.
…e task entry Until the anchor lands, a short grant is handed back to Spark whole. In the gap between that grant returning and its rollback release landing, a sibling acquire of the same task can enter Spark and park if another task has left the pool meanwhile, and the release then removes the task's entry under it. Carriers now run one at a time from the bridge acquire through the rollback release. Releases never take that lock, so a parked carrier cannot hold up the release it waits for, and Spark serializes a task's acquires anyway.
andygrove
added
enhancement
|
@sunchao the anchor redesign replacing the deferral is pushed, with the Spark-faithful stub and your P1 and late-arrival cases as tests. Ready for another look. |
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
The anchor invariant this design rests on does not hold during the window in which the anchor is being established, and bootstrap only covers half of that window. In native/core/src/execution/memory_pools/fair_pool.rs around line 345 every carrier takes bootstrap and holds it across the bridge acquire, so until the anchor lands the pool is back to one blocking JNI call at a time, which is the serialization this PR set out to remove. That lock also serializes only this pool's own carriers. CometTaskMemoryManager.internal is the task-wide TaskMemoryManager, and CometUnifiedShuffleMemoryAllocator is another off-heap consumer on the same task, so if it frees its last page while a pre-anchor carrier is parked inside Spark, memoryForTask loses the entry and the waiter hits the same NoSuchElementException the anchor exists to prevent. Could the anchor instead be taken once as a standalone one byte acquire when the pool is set up, so that no ordinary request ever carries it and Anchor::Requested, bootstrap, and the surplus path all go away?
The extra byte on the request path also turns grants Spark would have satisfied into failures. exact_fit_grant_without_the_anchor_is_rejected_as_short pins the case where Spark backs additional but not additional + 1 and the whole grant is handed back as short. grow at line 272 is try_grow(...).unwrap(), so a caller on that path gets a panic rather than a spill. It is a knife edge, but it is a knife edge that only shows up under memory pressure. Is that acceptable, or is it another argument for keeping the anchor off the request path entirely?
#5466 rewrites the same admission check in this file and moves the fair-share basis from the pool-wide state.used to reservation.size(). DataFusion does not update reservation.size() until try_grow returns, so that change would undercut the reasoning in the comment at line 307 that reserving under the lock is what makes the unlocked window safe. Whichever of the two lands second needs a deliberate rebase rather than a textual merge.
Worth noting that CI has not run on b7007ea9 at all, since the workflows are still sitting at action_required. The microbenchmarks sunchao asked for in the first round are also still open, so the only evidence behind the performance claim right now is the stub timing.
…t path CometFairMemoryPool::try_new acquires the one byte anchor from Spark before the pool is usable and fails construction if Spark declines it, so every live pool holds the anchor for its whole life and no ordinary request carries it. This removes Anchor::Requested, the bootstrap mutex that serialized pre-anchor acquires, the surplus release, and the exact-fit rejection; try_grow asks for exactly the requested bytes again. A sibling consumer of the same task freeing its last page can no longer drop the task's memoryForTask entry under a parked acquire, since the anchor keeps the balance above zero. Because construction can now park inside Spark, acquire_task_shared_pool runs the create closure without the process-wide registry lock and keeps the first pool registered when two plans of one task create at once. create_memory_pool returns CometResult and createPlan propagates the error.
|
@andygrove the anchor is now a standalone one byte acquire at pool setup (d004f04). One consequence worth calling out: construction can now park inside Spark, and On On #5466: it moves the admission basis from |
|
@sunchao the microbenchmarks from your first pass, against a real Spark task memory manager rather than the stub: the harness launches a JVM in-process through the jni crate's invocation feature, builds a Spark 4.1.3
The release path is where the change shows: at 4 and 8 threads its p99 drops from tens of microseconds to single digits because a release no longer waits behind an in-flight acquire. Grow p99 is unchanged, since that is Spark's own work under the call, while mean grow latency and throughput improve 1.5x to 2.5x under contention. Head sees fewer short grants in the partial scenarios because it cycles past the release faster. Registration churn counts only depend on run length. The harness needs a JDK and a full Spark classpath so it is not a criterion bench in tree; the source and raw output are available if wanted. macOS arm64, single machine. |
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
This is the shape I was hoping for, and it is a lot simpler than what it replaces. Anchor::Requested, the bootstrap mutex, the surplus release and the exact-fit rejection are all gone, and try_grow is back to reserve-under-lock, call-unlocked, roll-back-on-shortfall with no extra byte riding along. Moving the create step out from under the process-wide registry lock was the right call and I would not have thought to check it until it bit someone: with the anchor acquire able to park, holding TASK_SHARED_MEMORY_POOLS across it would have stalled plan creation for every task on the executor. The std::ptr::eq guard in TaskSharedMemoryPool::drop already covers the loser, so a concurrent create hands its anchor back without deregistering the winner. Thanks also for running the benchmarks against a real TaskMemoryManager rather than the stub; the release p99 collapsing from tens of microseconds to single digits at 4 and 8 threads is the number that actually demonstrates the claim, and it is the one the stub could not show.
One thing came out of the move that I do not think we should ship as is.
A declined anchor now fails the task instead of spilling
with_bridge returns CometError::Internal when Spark grants less than the byte, and createPlan propagates it, so the task dies before any work is attempted. Walking Spark's ExecutionMemoryPool.acquireMemory with numBytes = 1: maxToGrant = min(1, max(0, maxMemoryPerTask - curMem)), so a task whose curMem has already reached maxMemoryPerTask gets maxToGrant = 0, and the park guard curMem + toGrant < minMemoryPerTask is false for it because minMemoryPerTask is half of maxMemoryPerTask. It does not wait. It returns 0, and the pool refuses to exist.
That is reachable without anything exotic. The task only needs to be holding its share of execution memory at the moment the first native plan is created, and CometTaskMemoryManager.internal is the task-wide manager, so CometUnifiedShuffleMemoryAllocator and any Spark-side operator in a mixed plan count toward curMem. The concrete sequence I would worry about is a task that runs one native plan, closes it so the shared pool drops, allocates JVM shuffle pages, and then creates a second native plan: lookup misses, try_new runs again, and Spark declines. On main that situation produces a short grant on the first real try_grow and DataFusion spills. On this branch it is a task failure with no spill available as an escape, and the message reads like an internal error rather than memory pressure.
Could a declined anchor be non-fatal instead? Keep a flag for "anchor not held" and re-attempt a standalone one byte acquire at the top of try_grow while it is unset. That keeps everything this commit won: the anchor never rides on a real request, so there is no exact-fit hazard, no surplus to hand back and no bootstrap mutex, and the request path stays plain. The cost is one extra JNI call per try_grow only while the anchor is missing, which is the rare path by construction. The window where the pool holds reservation bytes without the anchor is then real but bounded, and it is strictly better than the window on main, which is the whole first acquire. A test for the declined-at-setup case would pin it either way; right now the stub only covers the grant.
The parking case reads fine to me and I convinced myself it cannot deadlock: a task parked in try_new holds nothing, so a pool fully consumed by tasks that are all parked in construction is not a reachable state, and maybeGrowPool covers the storage-borrow side.
Two smaller things. with_bridge's short-grant branch returns without releasing granted, which is dead today because ANCHOR_BYTES is 1 so a short grant is necessarily 0, but it would leak silently if that constant ever moved; a debug_assert or a one-line comment tying the two together would be cheap. And bridge.acquire at construction is the one bridge call not wrapped in catch_unwind. That is correct as written, since there is no optimistic reservation to roll back and try_unwrap_or_throw catches the panic at the JNI boundary, but the comment above the catch_unwind in try_grow explains why it is needed there without saying why it is not needed here.
Your read on #5466 matches mine, including that the per-reservation basis has to be applied to the reserve-under-lock step rather than replacing it, and that the two comments explaining the state.used basis need rewriting rather than merging. Worth agreeing an order with @peterxcli so whoever is second plans for it.
Spark grants zero bytes to a task already at its share, which a task holding JVM shuffle pages when its second native plan is created can reach. Failing pool construction there fails the task with no spill available, where main produces a short grant and spills. Create the pool without the anchor when Spark declines it, and have each grow retry a standalone one byte acquire first while it is missing, so the anchor never rides on a real request and the task spills on the short grant instead. A second grow racing the retry hands its byte back; drop releases the anchor only when held. A compile-time assert ties the zero grant reasoning to the one byte anchor size.
Done that way. The short-grant branch in Tests: the declined-at-setup case now asserts the pool exists holding nothing, reports a short grant while the sibling still holds the share, takes the byte on the first grow after the share frees, makes exactly one JVM call per grow from then on, and returns the byte at drop; a second test pins that a pool that never held the anchor releases nothing on drop; a third drives two retries into the stub's acquire gate before either wins and checks one anchor is kept and one byte returned. Thirty loops in debug and release, no failures. On #5466: left a note there proposing that whoever lands second rebases, with the per-reservation basis applied to the reserve-under-lock step rather than replacing it. |
3 participants
Which issue does this PR close?
No dedicated issue. Adjacent to #5494, which made task-shared pools the norm and so widened the blast radius of this lock.
Rationale for this change
CometFairMemoryPoolheld its internal mutex across the JNI acquire and release calls into Spark'sTaskMemoryManager. That call can block for a long time while Spark spills other consumers, and while it blocked, every other native thread sharing the task pool sat behind the lock, including plain releases that needed nothing from the JVM. The sibling unified pool already avoids this.What changes are included in this PR?
The admission check and the reservation now happen as one short locked step (the fair limit couples used bytes and the consumer count, so this part genuinely needs mutual exclusion), then the blocking JVM call runs with no lock held, and the reservation rolls back if the JVM declines, grants partially, or the call panics. Fairness semantics are unchanged: concurrent grows still cannot jointly exceed pool_size divided by the consumer count, and registering a new consumer still only blocks further growth rather than clawing back existing reservations. Going fully lock-free like the unified pool was considered and rejected, since separate atomics would let a register or unregister slip between reading the count and committing the reservation.
Spark's
ExecutionMemoryPoolremoves a task's accounting entry the moment its balance reaches zero, and an acquire parked inside Spark indexes that entry when it wakes, so a release that zeroes the balance while another acquire is parked crashes the waiter. With the lock no longer held across the call that window is real, so the pool takes a one byte anchor once at setup, before it exists, and returns it when it drops; no request ever carries it and the balance never reaches zero mid task. A parked one-byte acquire at setup blocks plan creation until memory frees, which is Spark's normal behavior for a starved task. Spark declines the byte outright, with a zero grant, for a task already at its share, which a task holding JVM shuffle pages when its second native plan is created can reach; the pool then exists without its anchor and each grow retries the byte as a request of its own until it is held, so the task spills on the short grant instead of failing at plan creation. That costs one extra JNI call per grow only while the anchor is missing, and the window without an anchor is bounded and strictly shorter than the whole first acquire on main. Because construction can now park, the task-shared pool registry creates the pool outside its process-wide lock and re-checks on insert; a concurrent loser drops its own pool and returns its anchor after the lock is released.Two side effects worth naming. The JNI boundary moved behind a small internal trait so the pool can be tested without a live JVM (neither pool had any tests before). And the old code could deadlock if a blocked acquire ever re-entered the pool on the same thread via a spill callback, since the lock was held across the call; that hazard is gone by construction.
How are these changes tested?
Twenty-five pool tests and nine registry tests, all new. The stub models Spark 4.1.3's
ExecutionMemoryPool: the per-task entry lifecycle, the 1/N and 1/(2N) shares, the wait loop, andnotifyAllon release, with a bounded wait that fails a test instead of hanging. Covered: fairness rejection without reaching Spark, limit tightening on register, partial-grant rollback with the excess returned, acquire failure and panic rollback, zero-size no-op, over-shrink panic, a parked acquire not stalling a concurrent release, an acquire below its minimum share being granted after the holder frees memory in full, a release already on its way when a late acquire parks, a sibling consumer of the same task freeing its last page while an acquire is parked, a second acquire not queuing behind a parked first one, an exact-fit first grant being accepted, the anchor taken at setup and returned once at drop, construction failing cleanly when the bridge errors on the anchor, a declined anchor leaving a usable pool that reports short grants and takes the byte on the first grow after the share frees with no second request once held, two grows retrying the missing anchor at once keeping exactly one byte, the registry creating outside its lock and keeping exactly one anchor when two creates race, and an eight-thread stress test asserting the accounting never exceeds the fair limit and nets to zero. All looped 30 times in debug and release with no failures. Full core crate suite passes, clippy with warnings denied and fmt are clean.Microbenchmarks against a real Spark 4.1.3
TaskMemoryManagerover an off-heapUnifiedMemoryManager, driven through the JNI bridge with base and head in one binary (details and the full table in the discussion): at 4 and 8 native threads the release p99 falls from tens of microseconds to single digits because a release no longer waits behind an in-flight acquire, mean grow latency and throughput improve 1.5x to 2.5x under contention, grow p99 is unchanged since that is Spark's own work, and errors and final balances are zero on every row.